import pickle


class Sequence:
    def __init__(self, prompt_tokens):
        # 初始化所有属性（确保每个属性都有定义）
        self.token_ids = prompt_tokens  # 完整token序列
        self.num_prompt_tokens = len(prompt_tokens)  # prompt长度
        self.num_completion_tokens = 0  # 生成的回答长度（初始为0）
        self.num_tokens = self.num_prompt_tokens  # 总长度 = prompt长度（初始无回答）
        self.num_cached_tokens = 0  # 已缓存的token数
        self.block_table = []  # KV缓存块表
        self.last_token = self.token_ids[-1] if self.token_ids else None  # 初始最后一个token是prompt的最后一个

    def generate_token(self, new_token):
        """模拟生成一个新token，更新所有关联属性"""
        self.token_ids.append(new_token)
        self.num_completion_tokens += 1  # 回答长度+1
        self.num_tokens += 1  # 总长度+1
        self.last_token = new_token  # 更新最后一个token

    def __getstate__(self):
        """自定义序列化：只保存必要信息，节省空间"""
        return (
            self.num_tokens,
            self.num_prompt_tokens,
            self.num_cached_tokens,
            self.block_table,
            # 未生成回答时保存完整token_ids，已生成则只保存最后一个token
            self.token_ids if self.num_completion_tokens == 0 else self.last_token
        )

    def __setstate__(self, state):
        """自定义反序列化：先补全所有属性，再判断逻辑"""
        # 步骤1：恢复state中前4个基础属性
        self.num_tokens, self.num_prompt_tokens, self.num_cached_tokens, self.block_table = state[:-1]

        # 关键修复：先计算初始化 num_completion_tokens（总token数 - prompttoken数）
        self.num_completion_tokens = self.num_tokens - self.num_prompt_tokens

        # 步骤2：根据 num_completion_tokens 恢复token相关属性
        if self.num_completion_tokens == 0:
            # 未生成回答：直接恢复完整token_ids，更新last_token
            self.token_ids = state[-1]
            self.last_token = self.token_ids[-1] if self.token_ids else None
        else:
            # 已生成回答：恢复last_token，token_ids用占位符+最后一个token填充
            self.last_token = state[-1]
            self.token_ids = [None] * self.num_tokens  # 占位（实际场景可按需优化）
            self.token_ids[-1] = self.last_token  # 只保证最后一个token有效

    def __repr__(self):
        """方便打印查看状态"""
        return (f"Sequence(\n"
                f"  num_tokens={self.num_tokens}, \n"
                f"  num_prompt_tokens={self.num_prompt_tokens}, \n"
                f"  num_completion_tokens={self.num_completion_tokens}, \n"
                f"  last_token={self.last_token}, \n"
                f"  token_ids={self.token_ids[:5]}...)"  # 只显示前5个token，避免过长
                )


# ------------------- 测试验证 -------------------
if __name__ == "__main__":
    # 1. 测试「未生成回答」的序列（仅含prompt）
    print("=" * 50)
    prompt_tokens = [101, 102, 103, 104]  # 模拟prompt的token ID
    seq1 = Sequence(prompt_tokens)
    print("原始序列（未生成回答）：")
    print(seq1)

    # 序列化
    pickled_seq1 = pickle.dumps(seq1)
    print(f"\n序列化后数据长度：{len(pickled_seq1)} 字节")

    # 反序列化
    unpickled_seq1 = pickle.loads(pickled_seq1)
    print("\n反序列化后序列（未生成回答）：")
    print(unpickled_seq1)
    print("反序列化后token_ids是否完整：", unpickled_seq1.token_ids == prompt_tokens)  # 应该为True

    # 2. 测试「已生成回答」的序列
    print("\n" + "=" * 50)
    seq2 = Sequence(prompt_tokens)
    seq2.generate_token(105)  # 生成第1个回答token
    seq2.generate_token(106)  # 生成第2个回答token
    print("原始序列（已生成2个回答token）：")
    print(seq2)

    # 序列化（此时只保存last_token，数据更短）
    pickled_seq2 = pickle.dumps(seq2)
    print(f"\n序列化后数据长度：{len(pickled_seq2)} 字节")  # 比seq1短

    # 反序列化
    unpickled_seq2 = pickle.loads(pickled_seq2)
    print("\n反序列化后序列（已生成回答）：")
    print(unpickled_seq2)
    print("反序列化后last_token是否正确：", unpickled_seq2.last_token == 106)  # 应该为True